# 准备量化配置
model.qconfig = torch.quantization.get_default_qconfig('fbgemm')

# 准备量化
torch.quantization.prepare(model, inplace=True)

# 转换模型为量化模型
torch.quantization.convert(model, inplace=True)


class BinaryQuantize(nn.Module):
    def forward(self, x):
        return x.sign()


    self.binary_quantize = BinaryQuantize()
    x = self.binary_quantize(x)


class TernaryQuantize(nn.Module):
    def __init__(self, threshold=0.5):
        super(TernaryQuantize, self).__init__()
        self.threshold = threshold

    def forward(self, x):
        return torch.where(x.abs() > self.threshold, torch.sign(x), torch.zeros_like(x))


    self.ternary_quantize = TernaryQuantize()
    x = self.ternary_quantize(x)


# 使用自动混合精度训练
with autocast():
    outputs = model(inputs)
    loss = nn.MSELoss()(outputs, targets)

scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()
